{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# tesseract OCR\n",
    "\n",
    "from PIL import Image\n",
    "import pytesseract\n",
    "\n",
    "def img_to_str_tesseract(image_path, lang='chi_sim'):\n",
    "    return pytesseract.image_to_string(Image.open(image_path), lang)\n",
    "\n",
    "print(img_to_str_tesseract('example.png'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 百度 OCR\n",
    "\n",
    "from aip import AipOcr\n",
    "\n",
    "config = {\n",
    "    'appId': '',\n",
    "    'apiKey': '',\n",
    "    'secretKey': ''\n",
    "}\n",
    "client = AipOcr(**config)\n",
    "\n",
    "def img_to_str_baidu(image_path):\n",
    "    with open(image_path, 'rb') as fp:\n",
    "        image = fp.read()\n",
    "        result = client.basicGeneral(image)\n",
    "        if 'words_result' in result:\n",
    "            return '\\n'.join([w['words'] for w in result['words_result']])\n",
    "    return \"\"\n",
    "\n",
    "print(img_to_str_baidu('example.png'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import fitz\n",
    "import time\n",
    "import re\n",
    "import os\n",
    "import sys\n",
    "import numpy as np\n",
    "\n",
    "TMPDIR = 'tmp/'\n",
    "PARSEIMG = True\n",
    "OCR_ONLINE = False\n",
    "\n",
    "# 去掉文中多余的回车\n",
    "def adjust(inpath, outpath):\n",
    "    f = open(inpath)\n",
    "    lines = f.readlines()\n",
    "    arr = [len(line) for line in lines]\n",
    "    length = np.median(arr) # 行字符数中值\n",
    "    \n",
    "    string = \"\"\n",
    "    for line in lines:\n",
    "        if len(line) >= length and line[-1]=='\\n':\n",
    "            string += line[:-1] # 去掉句尾的回车\n",
    "        elif line == '-----------\\n':\n",
    "            pass\n",
    "        else:\n",
    "            string += line\n",
    "    write_file(outpath, string, 'w')\n",
    "    return\n",
    "\n",
    "# 写入文件\n",
    "def write_file(path, text, ftype, debug=False):\n",
    "    with open(path, ftype) as f:\n",
    "        if debug:\n",
    "            print(\"write\", len(text))\n",
    "        f.write(text)\n",
    "        f.close()\n",
    "\n",
    "# 删除文件  \n",
    "def remove(path):\n",
    "    if not os.path.exists(path):\n",
    "        return\n",
    "    if os.path.isfile(path):\n",
    "        os.remove(path)\n",
    "        return\n",
    "    dirs = os.listdir(path)\n",
    "    for f in dirs:\n",
    "        file_name = os.path.join(path, f)\n",
    "        if os.path.isfile(file_name):\n",
    "            os.remove(file_name)\n",
    "        else:\n",
    "            remove(file_name)\n",
    "    os.rmdir(path)\n",
    "\n",
    "# 解析PDF文件\n",
    "def parse(inpath, outpath):\n",
    "    remove(TMPDIR) # 清除临时目录 \n",
    "    os.mkdir(TMPDIR)\n",
    "    remove(outpath) # 清除输出文件\n",
    "\n",
    "    t0 = time.clock()\n",
    "    doc = fitz.open(inpath)\n",
    "    lenXREF = doc.xrefLength()\n",
    "    print(\"文件名:{}, 页数: {}, 对象: {}\".format(inpath, len(doc), lenXREF - 1))\n",
    "\n",
    "    imgcount = 0\n",
    "    for i,page in enumerate(doc):\n",
    "        t1 = time.clock()\n",
    "        # 文字\n",
    "        text = page.get_text()\n",
    "        if len(text) > 0:\n",
    "             write_file(outpath, text, 'a')\n",
    "        # 图片        \n",
    "        imglist = page.get_images() # 解析该页中图片\n",
    "        for item in imglist:\n",
    "            xref = item[0]\n",
    "            pix = fitz.Pixmap(doc, xref)\n",
    "            new_name = \"{}.png\".format(imgcount)\n",
    "            # 如果pix.n<5,可以直接存为PNG\n",
    "            path = os.path.join(TMPDIR, new_name)\n",
    "            if pix.n < 5:\n",
    "                pix.writePNG(path)\n",
    "            # 否则先转换CMYK\n",
    "            else:\n",
    "                pix0 = fitz.Pixmap(fitz.csRGB, pix)\n",
    "                pix0.writePNG(path)\n",
    "                pix0 = None\n",
    "            pix = None\n",
    "            if OCR_ONLINE:\n",
    "                text = img_to_str_baidu(path)\n",
    "            else:\n",
    "                text = img_to_str_tesseract(path)\n",
    "            print(\"img->text\", text)\n",
    "            write_file(outpath, text, 'a')\n",
    "            write_file(outpath, '\\n' + '-----------' + '\\n', 'a')\n",
    "            imgcount += 1\n",
    "        print(\"page {} 运行时间:{}s\".format(i, {t1 - t0}))\n",
    "\n",
    "if __name__ == '__main__':\n",
    "    pdffile = 'xxxxx'\n",
    "    tmpfile = pdffile.replace('pdf','tmp')\n",
    "    txtfile = pdffile.replace('pdf','txt')\n",
    "    parse(pdffile, tmpfile)\n",
    "    adjust(tmpfile, txtfile)\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
